# Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
import random
from typing import Union, List

from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator
from atk.configs.case_config import InputCaseConfig, CaseConfig


@GENERATOR_REGISTRY.register("ascend_generate_expand_into_jagged_permute")
class ExpandIntoJaggedPermuteGenerator(CaseGenerator):

    def __init__(self, config):
        super().__init__(config)


    def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
        #input_offsets和output_offsets比permute长度多1
        case_config.inputs[1].shape = [case_config.inputs[0].shape[0] + 1]
        case_config.inputs[2].shape = [case_config.inputs[0].shape[0] + 1]
        return case_config